package com.peng.leetcode.trie;

import java.util.*;

/**
 * WordSize
 *
 * 使用字典树 + 小顶堆统计词频
 *
 * @author: lupeng6
 * @create: 2021/4/20 16:19
 */
public class WordSize {


    public static void main(String[] args) {

    }

    // 小顶堆统计 top
    private PriorityQueue<TopNode> topNodes = new PriorityQueue<>(Comparator.comparingInt(n -> n.size));

    //字典树的root节点
    private Node root = new Node();


    /**
     * 放入 word， 更新size 和小顶堆
     *
     * @author lupeng6
     * @date 2021/4/20 16:32
     */
    public void put(String word) {
        int len = word.length();
        Node curr = root;
        for (int i = 0; i < len; i++) {
            char c = word.charAt(i);
            curr.child.putIfAbsent(c, new Node());
            curr = curr.child.get(c);
        }
        curr.size++;
        curr.end = true;

        // 更新小顶堆
        if (topNodes.peek().size < curr.size) {
            topNodes.poll();
            topNodes.offer(new TopNode(curr.size, word));
        }
    }

    /**
     * 获取 word 出现的频率
     *
     * @author lupeng6
     * @date 2021/4/20 16:32
     */
    public int getSize(String word) {
        int len = word.length();
        Node curr = root;
        for (int i = 0; i < len; i++) {
            char c = word.charAt(i);
            curr = curr.child.get(c);
            if (curr == null) {
                return 0;
            }
        }
        if (!curr.end) {
            return 0;
        }
        return curr.size;
    }

    /**
     * 获取 出现次数最高的 前100个 word
     *
     * @author lupeng6
     * @date 2021/4/20 16:32
     */
    public List<String> getTop100() {
        List<String> ans = new ArrayList<>(topNodes.size());
        Iterator<TopNode> iterator = topNodes.iterator();
        while (iterator.hasNext()) {
            ans.add(iterator.next().word);
        }
        return ans;
    }



    static class Node {
        Map<Character, Node> child = new HashMap<>();
        boolean end = false;
        int size = 0;
    }

    static class TopNode {
        int size;
        String word;

        public TopNode(int size, String word) {
            this.size = size;
            this.word = word;
        }

        public TopNode() {
        }
    }
}
